import numpy as np
import torch
import torch.nn as nn


def get_hop_distance(num_node, edge, max_hop=1):
    A = np.zeros((num_node, num_node))
    for i, j in edge:
        A[j, i] = 1
        A[i, j] = 1

    # compute hop steps
    hop_dis = np.zeros((num_node, num_node)) + np.inf
    transfer_mat = [np.linalg.matrix_power(A, d) for d in range(max_hop + 1)]
    arrive_mat = (np.stack(transfer_mat) > 0)
    for d in range(max_hop, -1, -1):
        hop_dis[arrive_mat[d]] = d
    return hop_dis


def normalize_digraph(A):
    Dl = np.sum(A, 0)
    num_node = A.shape[0]
    Dn = np.zeros((num_node, num_node))
    for i in range(num_node):
        if Dl[i] > 0:
            Dn[i, i] = Dl[i] ** (-1)
    AD = np.dot(A, Dn)
    return AD


class Graph:
    def __init__(self, layout='openpose', strategy='uniform', max_hop=1, dilation=1, num_nodes=15):
        self.max_hop = max_hop
        self.dilation = dilation
        self.num_node = num_nodes

        self.get_edge(layout)
        self.hop_dis = get_hop_distance(self.num_node, self.edge, max_hop=max_hop)
        self.get_adjacency(strategy)

    def __str__(self):
        return str(self.A)

    def get_edge(self, layout):
        if layout == 'openpose':
            self_link = [(i, i) for i in range(self.num_node)]
            neighbor_link = [(0, 1), (1, 2), (1, 3), (2, 4), (1, 5), (3, 6), (4, 7), (5, 8), (6, 10), (5, 9), (9, 11),
                             (10, 12), (12, 13), (13, 14)]
            self.edge = self_link + neighbor_link
            self.center = 1
        else:
            raise ValueError("Do Not Exist This Layout.")

    def get_adjacency(self, strategy):
        valid_hop = range(0, self.max_hop + 1, self.dilation)
        adjacency = np.zeros((self.num_node, self.num_node))
        for hop in valid_hop:
            adjacency[self.hop_dis == hop] = 1
        normalize_adjacency = normalize_digraph(adjacency)

        if strategy == 'uniform':
            A = np.zeros((1, self.num_node, self.num_node))
            A[0] = normalize_adjacency
            self.A = A
        elif strategy == 'three_streams':
            self_loop = np.eye(self.num_node)
            centripetal = np.zeros((self.num_node, self.num_node))
            centrifugal = np.zeros((self.num_node, self.num_node))

            for i in range(self.num_node):
                for j in range(self.num_node):
                    if self.hop_dis[i, j] == 1:
                        if i == self.center:
                            centrifugal[i, j] = 1
                        elif j == self.center:
                            centripetal[i, j] = 1

            self_loop = normalize_digraph(self_loop)
            centripetal = normalize_digraph(centripetal)
            centrifugal = normalize_digraph(centrifugal)

            A = np.stack([self_loop, centripetal, centrifugal])
            self.A = A
        elif strategy == 'adaptive':
            # 这里只是简单示例，实际中可能需要更复杂的计算
            self.A = nn.Parameter(torch.eye(self.num_node).unsqueeze(0), requires_grad=True)
        else:
            raise ValueError("Do Not Exist This Strategy")
